TOP = __name__ == '__main__'

from IPython.core.magic import  (Magics, magics_class, cell_magic, line_magic, line_cell_magic)

class B(dict):

    def __getattr__(self, name):

        return self.__getitem__(name)

@magics_class

class MyMagics(Magics):

    def __init__(self, shell=None,  **kwargs):

        super().__init__(shell=shell, **kwargs)

        self._vars = B()

        shell.user_ns['VARS'] = self._vars

​

    @cell_magic

    def var(self, line, cell):

        self._vars[line.strip()] = cell.strip()

    @line_cell_magic

    def top(self, line, cell=None):

        if TOP:

            if cell is None:

                cell = line

            ip = get_ipython()

            res = ip.run_cell(cell)

get_ipython().register_magics(MyMagics)

import sys

%top assert sys.version_info[0:2] == (3, 7)

from subprocess import run

import os

def do(command, env=None, shell=False, log=False, **args):

    result = run(command, universal_newlines=True, shell=shell,

                  env=dict(os.environ, **({} if env is None else env)),

                 capture_output=True, **args)

    if log:

        with open('build/do.log', 'a+') as f:

            print(json.dumps({'cmd':command, 'env':env, 'exitcode':result.returncode}), env, file=f)

    return result

import random

random.seed(0)

import shutil

%%top

if shutil.which('lsb_release'):

    res = do(['lsb_release', '-d']).stdout

elif shutil.which('sw_vers'):

    res = do(['sw_vers']).stdout

else:

    assert False

res

%top do(['jupyter', '--version']).stdout

def pip_install(v):

    return do(['pip', 'install', '-qqq', *v.split(' ')]).returncode

%top pip_install('fuzzingbook==0.8.0')

%top pip_install('astor graphviz scipy')

%top pip_install('jupyter_contrib_nbextensions jupyter_nbextensions_configurator')

%top do(['jupyter','contrib','nbextension','install','--user']).returncode

def nb_enable(v): return do(['jupyter','nbextension','enable',v]).returncode

%top do(['jupyter','nbextensions_configurator','enable','--user']).returncode

%top nb_enable('toc2/main')

%top nb_enable('collapsible_headings/main')

%top nb_enable('execute_time/ExecuteTime')

%top nb_enable('codefolding/main')

%top do(['rm', '-rf','build']).returncode

%%var Mimid

# [(

Testing Mimid

# )]

%top VARS['Mimid']

%%var calc_src↔​

%%var mathexpr_src↔​

%%var microjson_src↔​

%%var urlparse_src↔​

%%var netrc_src↔​

%%var cgidecode_src↔​

# [(

program_src = {

    'calculator.py': VARS['calc_src'],

    'mathexpr.py': VARS['mathexpr_src'],

    'urlparse.py': VARS['urlparse_src'],

    'netrc.py': VARS['netrc_src'],

    'cgidecode.py': VARS['cgidecode_src'],

    'microjson.py': VARS['microjson_src']

}

# )]

import ast

import astor

class InRewriter(ast.NodeTransformer):

    def wrap(self, node):

        return ast.Call(func=ast.Name(id='taint_wrap__', ctx=ast.Load()), args=[node], keywords=[])

class InRewriter(InRewriter):

    def visit_Compare(self, tree_node):

        left = tree_node.left

        if not tree_node.ops or not isinstance(tree_node.ops[0], ast.In):

            return tree_node

        mod_val = ast.Call(

            func=ast.Attribute(

                value=self.wrap(left),

                attr='in_'),

            args=tree_node.comparators,

            keywords=[])

        return mod_val

def rewrite_in(src):

    v = ast.fix_missing_locations(InRewriter().visit(ast.parse(src)))

    source = astor.to_source(v)

    return "%s" % source

from fuzzingbook.fuzzingbook_utils import print_content

%top print_content(rewrite_in('s in ["a", "b", "c"]'))

class Rewriter(InRewriter):

    def init_counters(self):

        self.if_counter = 0

        self.while_counter = 0

methods = []

Epsilon = '-'

NoEpsilon = '='

class Rewriter(Rewriter):

    def wrap_in_method(self, body, args):

        method_name_expr = ast.Str(methods[-1])

        my_args = ast.List(args.args, ast.Load())

        args = [method_name_expr, my_args]

        scope_expr = ast.Call(func=ast.Name(id='method__', ctx=ast.Load()), args=args, keywords=[])

        return [ast.With(items=[ast.withitem(scope_expr, ast.Name(id='_method__'))], body=body)]

class Rewriter(Rewriter):

    def visit_FunctionDef(self, tree_node):

        self.init_counters()

        methods.append(tree_node.name)

        self.generic_visit(tree_node)

        tree_node.body = self.wrap_in_method(tree_node.body, tree_node.args)

        return tree_node

def rewrite_def(src):

    v = ast.fix_missing_locations(Rewriter().visit(ast.parse(src)))

    return astor.to_source(v)

%top print_content(rewrite_def('\n'.join(program_src['calculator.py'].split('\n')[12:19])), 'calculator.py')

class Rewriter(Rewriter):

    def wrap_in_outer(self, name, can_empty, counter, node):

        name_expr = ast.Str(name)

        can_empty_expr = ast.Str(can_empty)

        counter_expr = ast.Num(counter)

        method_id = ast.Name(id='_method__')

        args = [name_expr, counter_expr, method_id, can_empty_expr]

        scope_expr = ast.Call(func=ast.Name(id='stack__', ctx=ast.Load()),

                args=args, keywords=[])

        return ast.With(

            items=[ast.withitem(scope_expr, ast.Name(id='%s_%d_stack__' % (name, counter)))], 

            body=[node])

class Rewriter(Rewriter):

    def wrap_in_inner(self, name, counter, val, body):

        val_expr = ast.Num(val)

        stack_iter = ast.Name(id='%s_%d_stack__' % (name, counter))

        method_id = ast.Name(id='_method__')

        args = [val_expr, stack_iter, method_id]

        scope_expr = ast.Call(func=ast.Name(id='scope__', ctx=ast.Load()),

                args=args, keywords=[])

        return [ast.With(

            items=[ast.withitem(scope_expr, ast.Name(id='%s_%d_%d_scope__' % (name, counter, val)))], 

            body=body)]

class Rewriter(Rewriter):

    def process_if(self, tree_node, counter, val=None):

        if val is None: val = 0

        else: val += 1

        if_body = []

        self.generic_visit(tree_node.test)

        for node in tree_node.body: self.generic_visit(node)

        tree_node.body = self.wrap_in_inner('if', counter, val, tree_node.body)

​

        # else part.

        if len(tree_node.orelse) == 1 and isinstance(tree_node.orelse[0], ast.If):

            self.process_if(tree_node.orelse[0], counter, val)

        else:

            if tree_node.orelse:

                val += 1

                for node in tree_node.orelse: self.generic_visit(node)

                tree_node.orelse = self.wrap_in_inner('if', counter, val, tree_node.orelse)

class Rewriter(Rewriter):

    def visit_If(self, tree_node):

        self.if_counter += 1

        counter = self.if_counter

        #is it empty

        start = tree_node

        while start:

            if isinstance(start, ast.If):

                if not start.orelse:

                    start = None

                elif len(start.orelse) == 1:

                    start = start.orelse[0]

                else:

                    break

            else:

                break

        self.process_if(tree_node, counter=self.if_counter)

        can_empty = NoEpsilon if start else Epsilon  # NoEpsilon for + and Epsilon for *

        return self.wrap_in_outer('if', can_empty, counter, tree_node)

class Rewriter(Rewriter):

    def visit_While(self, tree_node):

        self.generic_visit(tree_node)

        self.while_counter += 1

        counter = self.while_counter

        test = tree_node.test

        body = tree_node.body

        assert not tree_node.orelse

        tree_node.body = self.wrap_in_inner('while', counter, 0, body)

        return self.wrap_in_outer('while', '?', counter, tree_node)

def rewrite_cf(src, original):

    v = ast.fix_missing_locations(Rewriter().visit(ast.parse(src)))

    return astor.to_source(v)

%top print_content('\n'.join(program_src['calculator.py'].split('\n')[12:19]), 'calculator.py')

%top print_content(rewrite_cf('\n'.join(program_src['calculator.py'].split('\n')[12:19]), 'calculator.py').strip(), filename='calculator.py')